#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
文生图HTTP路由
处理文字生成图片相关的HTTP接口
"""

import base64
import time
from flask import Blueprint, request, jsonify
from core import extensions
from services.streaming_tts_service import StreamingTTSService

# 创建蓝图
text_to_image_bp = Blueprint('text_to_image', __name__, url_prefix='/text_to_image')


@text_to_image_bp.route('/generate', methods=['POST'])
def generate_image_from_text():
    """
    文生图接口（增强版）
    根据文本描述生成图片，包含大模型改写、图片生成和语音播报
    
    请求参数:
        - prompt: 图片描述文本（必填）
        - num_images: 生成图片数量，1-5之间，默认1（可选）
        - size: 图片尺寸，支持 "1K"、"2K" 等，默认 "2K"（可选）
        - reference_images: 参考图片URL列表（可选）
        - watermark: 是否添加水印，默认 true（可选）
        - response_format: 返回格式，"url" 或 "b64_json"，默认 "url"（可选）
        - client_id: WebSocket客户端ID，用于语音播报（可选）
        - enable_llm_rewrite: 是否启用大模型改写提示词，默认 true（可选）
        - enable_tts: 是否启用语音播报，默认 true（可选）
    
    返回格式:
        {
            "code": 0,  // 0表示成功，其他表示失败
            "message": "成功生成 N 张图片",
            "data": {
                "images": [
                    {"url": "图片URL", "size": "图片尺寸"},
                    ...
                ],
                "count": 实际生成数量,
                "original_prompt": "原始描述",
                "rewritten_prompt": "改写后的描述",
                "size": "图片尺寸",
                "tts_played": true/false
            }
        }
    """
    try:
        # 检查服务是否可用
        if extensions.text_to_image_service is None:
            return jsonify({
                "code": 503,
                "message": "文生图服务未启用或初始化失败",
                "data": None
            }), 503
        
        # 获取请求参数
        data = request.get_json()
        if not data:
            return jsonify({
                "code": 400,
                "message": "请求参数不能为空",
                "data": None
            }), 400
        
        # 获取必填参数
        original_prompt = data.get('prompt', '').strip()
        if not original_prompt:
            return jsonify({
                "code": 400,
                "message": "图片描述（prompt）不能为空",
                "data": None
            }), 400
        
        # 获取可选参数
        num_images = data.get('num_images', 1)
        size = data.get('size', '2K')
        reference_images = data.get('reference_images', None)
        watermark = data.get('watermark', True)
        response_format = data.get('response_format', 'url')
        client_id = data.get('client_id', None)
        enable_llm_rewrite = data.get('enable_llm_rewrite', True)
        enable_tts = data.get('enable_tts', True)
        
        # 验证数量参数
        try:
            num_images = int(num_images)
            if num_images < 1:
                num_images = 1
            elif num_images > extensions.text_to_image_service.MAX_IMAGE_COUNT:
                return jsonify({
                    "code": 400,
                    "message": f"生成数量不能超过 {extensions.text_to_image_service.MAX_IMAGE_COUNT} 张",
                    "data": None
                }), 400
        except (ValueError, TypeError):
            return jsonify({
                "code": 400,
                "message": "生成数量必须是有效的整数",
                "data": None
            }), 400
        
        # ============= 步骤1：使用大模型改写提示词 =============
        rewritten_prompt = original_prompt
        
        if enable_llm_rewrite:
            try:
                print(f"🤖 [LLM改写] 开始改写提示词...")
                print(f"   原始提示词: {original_prompt}")
                print(f"   生成数量: {num_images}")
                
                # 构造大模型改写的提示
                llm_rewrite_prompt = f"""请将用户的图片生成需求改写成适合AI绘画的专业提示词。要求：

1. 如果用户要求生成多张图片（{num_images}张），请确保每张图片的描述连贯且有故事性
2. 包含详细的画面描述：主体、风格、色彩、光线、构图等
3. 使用适合AI绘画的专业术语
4. 如果是多张图片，用"生成{num_images}张"开头，然后用序号分别描述每张图片
5. 保持简洁，每张图片的描述控制在50字以内
6. 风格要统一协调

用户需求：{original_prompt}

请直接返回改写后的提示词，不要有其他说明文字。"""
                
                # 调用LLM客户端进行改写
                rewritten_result = extensions.llm_client.chat(
                    user_input=llm_rewrite_prompt,
                    system_prompt="你是一个AI绘画提示词专家，擅长将用户的简单描述改写成专业的AI绘画提示词。",
                    max_tokens=500,
                    temperature=0.7
                )
                
                # 清理改写结果（移除可能的前缀）
                rewritten_prompt = rewritten_result.strip()
                rewritten_prompt = rewritten_prompt.replace('改写后的提示词：', '').strip()
                rewritten_prompt = rewritten_prompt.replace('提示词：', '').strip()
                
                print(f"✅ [LLM改写] 改写完成")
                print(f"   改写后提示词: {rewritten_prompt}")
                
            except Exception as e:
                print(f"⚠️ [LLM改写] 改写失败，使用原始提示词: {str(e)}")
                rewritten_prompt = f"生成{num_images}张图片：{original_prompt}"
        else:
            print(f"ℹ️ [LLM改写] 已禁用，直接使用原始提示词")
            rewritten_prompt = f"生成{num_images}张图片：{original_prompt}"
        
        # 验证最终提示词
        is_valid, validation_msg = extensions.text_to_image_service.validate_prompt(rewritten_prompt)
        if not is_valid:
            return jsonify({
                "code": 400,
                "message": validation_msg,
                "data": None
            }), 400
        
        # ============= 步骤2：调用文生图服务生成图片 =============
        print(f"🎨 [图片生成] 开始生成图片...")
        
        result = extensions.text_to_image_service.generate_images(
            prompt=rewritten_prompt,
            num_images=num_images,
            size=size,
            reference_images=reference_images,
            watermark=watermark,
            response_format=response_format
        )
        
        # ============= 步骤3：语音播报生成结果 =============
        tts_played = False
        
        if result['success'] and enable_tts and client_id:
            try:
                print(f"🔊 [TTS播报] 开始语音播报...")
                
                # 构造播报文本
                tts_text = f"已成功生成 {result['data']['count']} 张图片，请在聊天框中查看。"
                
                # 通过SocketIO发送TTS请求
                tts_service = StreamingTTSService()
                
                # 生成TTS音频
                tts_result = tts_service.synthesize_speech(
                    text=tts_text,
                    voice_type='BV700_V2_streaming',
                    encoding='mp3',
                    speed_ratio=1.2,
                    volume_ratio=1.0,
                    pitch_ratio=1.0
                )
                
                if tts_result['success']:
                    # 通过WebSocket发送TTS音频
                    extensions.socketio.emit('streaming_tts_audio_chunk', {
                        'audio_data': tts_result['audio_base64'],
                        'chunk_index': 1,
                        'is_final': True
                    }, room=client_id)
                    
                    extensions.socketio.emit('streaming_tts_complete', {
                        'message': 'TTS播报完成',
                        'text': tts_text
                    }, room=client_id)
                    
                    tts_played = True
                    print(f"✅ [TTS播报] 播报成功")
                else:
                    print(f"⚠️ [TTS播报] 播报失败: {tts_result.get('message', '未知错误')}")
                    
            except Exception as e:
                print(f"⚠️ [TTS播报] 播报异常: {str(e)}")
        
        # ============= 构造响应数据 =============
        if result['success']:
            response_data = {
                "code": 0,
                "message": result['message'],
                "data": {
                    **result['data'],
                    "original_prompt": original_prompt,
                    "rewritten_prompt": rewritten_prompt,
                    "tts_played": tts_played,
                    "enable_llm_rewrite": enable_llm_rewrite
                }
            }
            
            print(f"✅ [完成] 图片生成流程全部完成")
            return jsonify(response_data)
        else:
            response_data = {
                "code": 500,
                "message": result['message'],
                "data": None
            }
            return jsonify(response_data), 500
        
    except Exception as e:
        print(f"❌ [错误] 图片生成失败: {str(e)}")
        response_data = {
            "code": 500,
            "message": f"图片生成失败: {str(e)}",
            "data": None
        }
        return jsonify(response_data), 500


@text_to_image_bp.route('/status', methods=['GET'])
def get_text_to_image_status():
    """
    获取文生图服务状态接口
    
    返回格式:
        {
            "code": 0,
            "message": "服务正常",
            "data": {
                "enabled": true/false,
                "max_images": 5,
                "supported_sizes": ["1K", "2K"],
                "supported_formats": ["url", "b64_json"]
            }
        }
    """
    try:
        if extensions.text_to_image_service is None:
            return jsonify({
                "code": 0,
                "message": "文生图服务未启用",
                "data": {
                    "enabled": False,
                    "max_images": 0,
                    "supported_sizes": [],
                    "supported_formats": []
                }
            })
        
        response_data = {
            "code": 0,
            "message": "文生图服务正常运行",
            "data": {
                "enabled": True,
                "max_images": extensions.text_to_image_service.MAX_IMAGE_COUNT,
                "model": extensions.text_to_image_service.model,
                "supported_sizes": ["1K", "2K"],
                "supported_formats": ["url", "b64_json"]
            }
        }
        return jsonify(response_data)
        
    except Exception as e:
        response_data = {
            "code": 500,
            "message": f"获取服务状态失败: {str(e)}",
            "data": None
        }
        return jsonify(response_data), 500

